import numpy as np
import torch
import torch.nn as nn
import torch.optim as optim
from torchvision import datasets, transforms
from torch.utils.data import DataLoader
import matplotlib.pyplot as plt
from matplotlib.widgets import Button, Slider
from PIL import Image
import cv2
import os

CNN_NETWORK = True
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"当前设备: {device}")

_transform = transforms.Compose([
    transforms.Grayscale(),
    transforms.ToTensor(),
    transforms.Resize((28, 28)),
    transforms.Lambda(lambda x: (x > 0.7).float()),
    transforms.Normalize((0.1307,), (0.3081,))
])  # 预处理器


# 1. 下载数据
def download_mnist():
    train_data = datasets.MNIST(root='./', download=True, train=True, transform=_transform)
    test_data = datasets.MNIST(root='./', download=True, train=False, transform=_transform)
    train_loader = DataLoader(dataset=train_data, batch_size=32, shuffle=True)
    test_loader = DataLoader(dataset=test_data, batch_size=32, shuffle=True)
    return train_loader, test_loader


# 2. 定义神经网络模型
class NetWork(nn.Module):
    def __init__(self):
        super(NetWork, self).__init__()
        self.model = nn.Sequential(
            nn.Linear(28 * 28, 64),
            nn.ReLU(),
            nn.Linear(64, 128),
            nn.ReLU(),
            nn.Linear(128, 64),
            nn.ReLU(),
            nn.Linear(64, 10),
            nn.LogSoftmax(dim=1)

        )

    def forward(self, x):
        x = x.view(-1, 28 * 28)
        return self.model(x)


class CNN(nn.Module):
    def __init__(self):
        super(CNN, self).__init__()
        self.conv1 = nn.Sequential(
            nn.Conv2d(1, 32, 3, 1, 1),  # 输入1通道，输出32通道，3x3卷积核，padding=1
            nn.BatchNorm2d(32),
            nn.ReLU(),
            nn.MaxPool2d(2))  # 池化层 2x2

        self.conv2 = nn.Sequential(
            nn.Conv2d(32, 64, 3, 1, 1),
            nn.BatchNorm2d(64),
            nn.ReLU(),
            nn.MaxPool2d(2))

        self.fc = nn.Sequential(
            nn.Linear(64 * 7 * 7, 128),  # 两次池化后28x28→14x14→7x7
            nn.ReLU(),
            nn.Dropout(0.5),
            nn.Linear(128, 10),
            nn.LogSoftmax(dim=1))

    def forward(self, x):
        x = self.conv1(x)
        x = self.conv2(x)
        x = x.view(x.size(0), -1)  # 展平多维卷积特征图
        x = self.fc(x)
        return x


# 3. 训练网络
def train_and_save():
    loss_func = nn.NLLLoss()
    _model = NetWork() if not CNN_NETWORK else CNN()
    train_loader, _ = download_mnist()
    epochs = 5
    for epoch in range(epochs):
        total_number = 0
        ture_number = 0
        for batch, (images, labels) in enumerate(train_loader):
            _model.zero_grad()
            out_puts = _model(images)
            loss = loss_func(out_puts, labels)
            loss.backward()
            optimization = optim.Adam(_model.parameters(), 0.001)
            optimization.step()
            # 统计良率
            for index, output in enumerate(out_puts):
                if torch.argmax(output).item() == labels[index]:
                    ture_number += 1
                total_number += 1
        print(f"epoch{(epoch + 1)}/{epochs}: {round(ture_number / total_number * 100, 3)}%")
    torch.save(_model.state_dict(), 'mnist.pth')
    return _model


# 4. 模型加载函数
def load_saved_model():
    print("Loading pre-trained model...")
    _model = NetWork() if not CNN_NETWORK else CNN()
    if os.path.exists('mnist.pth'):
        _model.load_state_dict(torch.load('mnist.pth', weights_only=True))
        _model.eval()
        print("Model loaded successfully!")
        return _model
    else:
        print("No saved model found. Training new model...")
        return train_and_save()


# 5. 测试模型精度
def test_model(model):
    total_number = 0
    ture_number = 0
    _, test_loader = download_mnist()
    for _, (images, labels) in enumerate(test_loader):
        outputs = model(images)
        for index, _ouput in enumerate(outputs):
            result = torch.argmax(_ouput).item()
            label = labels[index]
            ture_number = ture_number + 1 if result == label else ture_number
            total_number += 1
    print(f"Yield: {round(ture_number / total_number * 100, 3)}%")


# 6. 测试自定义的图片
def test_img(img_path):
    img = Image.open(img_path)
    img = _transform(img)
    print("img shape: ", img.shape)
    img = img.expand(1, -1, -1, -1)
    print("img shape: ", img.shape)
    _model = load_saved_model()
    outputs = _model(img)
    output = outputs[0]
    result = torch.argmax(output).item()
    print(f"Image result: {result}")


if __name__ == '__main__':
    download_mnist()
    model = load_saved_model()
    test_model(model)
    test_img("handwritten_digit.png")
